from mpl_finance import candlestick2_ohlc
from matplotlib import dates, ticker
from iplot import iplothelper
import datetime 



class DaliyResearch(object):
    """每日事件研究
    """
    def __init__(self, fig, ax, d_start, d_space, s, s1):
        """
        
        Arguments:
            fig {[type]} -- [description]
            ax {[type]} -- [description]
            s {[type]} -- [矩形开始时间 字符串)] '8:30'
            s1 {[type]} -- [矩形结束时间 字符串] '8:40'
            d_start {[int]} -- [图表开始时间 字符串] '2003-05-04 21:00:00'
            d_space {[type]} -- [description]
        """

        self.fig = fig
        self.ax = ax
        self.s = s
        self.s1 = s1
        self.d_start = iplothelper.str_to_datetime(d_start,'%Y-%m-%d %H:%M:%S')
        self.d_space = d_space
        self.draw()

    def draw(self):
        candleObj = iplothelper.getdata_by_timestamp(self.d_start, self.d_start + self.d_space)
        
        if len(candleObj.timestamp_list) <= 0:
            print("该index没有数据")
            return
        # 绘制基本K线图

        self.ax.clear()
        
        iplothelper.draw_candlestick_base2(self.fig, self.ax, candleObj)
       
        t = iplothelper.getZerotimestamp(self.d_start)
        d = datetime.datetime.strptime(self.s, '%H:%M:%S')
        d2 = datetime.datetime.strptime(self.s1, '%H:%M:%S')
        t1= t +  d.hour * 3600 + d.minute *60 + d.second
        t2 = t +  d2.hour * 3600 + d2.minute *60 + d2.second
        iplothelper.draw_rectangle(self.fig, self.ax, candleObj, t1, t2, 'Test')

    def press(self,event):
        print('press', event.key)

        if event.key == 'left':
            self.d_start -= 86400
        elif event.key == 'right':
            self.d_start += 86400
        self.draw()

        self.fig = fig
        self.ax = ax
        self.s = s
        self.s1 = s1
        self.d_start = iplothelper.str_to_datetime(d_start,'%Y-%m-%d %H:%M:%S')
        self.d_space = d_space
        self.draw()

    def draw(self):
        candleObj = iplothelper.getdata_by_timestamp(self.d_start, self.d_start + self.d_space)
        
        if len(candleObj.timestamp_list) <= 0:
            print("该index没有数据")
            return
        # 绘制基本K线图

        self.ax.clear()
        
        iplothelper.draw_candlestick_base2(self.fig, self.ax, candleObj)
       
        t = iplothelper.getZerotimestamp(self.d_start)
        d = datetime.datetime.strptime(self.s, '%H:%M:%S')
        d2 = datetime.datetime.strptime(self.s1, '%H:%M:%S')
        t1= t +  d.hour * 3600 + d.minute *60 + d.second
        t2 = t +  d2.hour * 3600 + d2.minute *60 + d2.second
        iplothelper.draw_rectangle(self.fig, self.ax, candleObj, t1, t2, 'Test')

    

class EventResearch(object):
    """基本面事件研究绘图板
    """

    def __init__(self, fig, ax, contry, ename, e_space, d_space):
        """初始化
        
        Arguments:
            ax {[type]} -- [matplotlib绘图对象]
            contry {[type]} -- [国家缩写]
            ename {[type]} -- [基本面事件名称]
            e_space {[type]} -- [基本面事件影响范围(秒)]
            d_space {[type]} -- [要观察的时间范围(秒)]
        """
        self.fig = fig
        self.ax = ax
        self.ename = ename
        self.contry = contry
        # 默认读取下标为0
        self.index = -2
        if d_space < e_space:
            print("观察时间范围必须大于基本面影响范围")
            return
        self.e_space = e_space
        self.d_space = d_space
        self.high_val = 0
        self.low_val = 0
        self.draw()
        
    def draw(self):
        eObj_list = iplothelper.get_eventlist_byName(self.contry, self.ename)
        if len(eObj_list) <= 0:
            print('未找到事件列表')
            return
        eObj = eObj_list[self.index]
        
        # 从当天00:00开始绘制
        d_start = iplothelper.getZerotimestamp(eObj.timestamp)
        d_end = d_start + self.d_space
        candleObj = iplothelper.getdata_by_timestamp(d_start,d_end)
        if len(candleObj.timestamp_list) <= 0:
            print("该index没有数据")
            return
        # 绘制基本K线图

        self.ax.clear()
        
        iplothelper.draw_candlestick_base2(self.fig, self.ax, candleObj)
        
        
        
        t = eObj.timestamp
        t1 = eObj.timestamp + self.e_space
        self.high_val,self.low_val = iplothelper.draw_rectangle(self.fig, self.ax, candleObj, t, t1, eObj.eventName)
        

    def press(self,event):
        print('press', event.key)

        if event.key == 'left':
            self.index -= 1
        elif event.key == 'right':
            self.index += 1
        self.draw()

class Maxdensity(object):
    def __init__(self, ax,fig, start_timestamp,end_timestamp,rs_stamp,re_stamp):
        self.ax = ax
        self.fig = fig
        ax.set_title('srolltest')
        
        # start_timestamp =  iplothelper.dateStr2timestamp(start_date,'%Y-%m-%d')
        # end_timestamp = iplothelper.dateStr2timestamp(end_date,'%Y-%m-%d')
        # rs_stamp =  iplothelper.dateStr2timestamp(range_statr_date,'%Y-%m-%d')
        # re_stamp =  iplothelper.dateStr2timestamp(range_end_date,'%Y-%m-%d')
        # 设置两边的时间戳


        cObj = iplothelper.getdata_by_timestamp(start_timestamp,end_timestamp)
        # 分割线
        self.division = 0
        for index,ts in enumerate(cObj.timestamp_list):
            # print(ts,end_timestamp)
            if ts >= re_stamp:
                self.division = index
                break

        r_cObj = iplothelper.getdata_by_timestamp(rs_stamp, re_stamp)

        self.max_denstity = iplothelper.fin_max_density_line(r_cObj)
        self.denstity_list = iplothelper.find_density_list(r_cObj)

        self.draw_candlestick(cObj)

    def draw_candlestick(self,cObj):
        
        datetime_list = cObj.dateStr_list
        open_list = cObj.open_list
        high_list = cObj.high_list
        low_list = cObj.low_list
        close_list = cObj.close_list
        self.ax.clear()
        candlestick2_ohlc(self.ax, open_list, high_list, low_list, close_list, width=0.6, colorup='red', colordown='black')

        def mydate(x, pos):
            try:
                return datetime_list[int(x)]
            except IndexError:
                return ''

        self.ax.xaxis.set_major_formatter(ticker.FuncFormatter(mydate))

        self.update()


    def update(self):
        # self.im.set_data(self.X[:, :, self.ind])
        # ax.set_ylabel('slice %s' % self.ind)

        self.fig.canvas.draw()



class IndexTracker(object):
    def __init__(self, ax,fig, start_date,space):
        self.ax = ax
        self.fig = fig
        ax.set_title('srolltest')
        
        start_timestamp =  iplothelper.dateStr2timestamp(start_date,'%Y-%m-%d')
        end_timestamp = start_timestamp + space
        # 设置两边的时间戳
        
        self.start_timestamp = start_timestamp
        
        self.end_timestamp = end_timestamp


        cObj = iplothelper.getdata_by_timestamp(start_timestamp,end_timestamp)

        self.draw_candlestick(cObj)

    def draw_candlestick(self,cObj):
        
        datetime_list = cObj.dateStr_list
        open_list = cObj.open_list
        high_list = cObj.high_list
        low_list = cObj.low_list
        close_list = cObj.close_list
        self.ax.clear()
        candlestick2_ohlc(self.ax, open_list, high_list, low_list, close_list, width=0.6, colorup='red', colordown='black')

        def mydate(x, pos):
            try:
                return datetime_list[int(x)]
            except IndexError:
                return ''

        self.ax.xaxis.set_major_formatter(ticker.FuncFormatter(mydate))

        self.update()

    def press(self,event):
        print('press', event.key)

        sql = ""

        if event.key == 'left':
            self.end_timestamp -= 86400
            self.start_timestamp -= 86400
            print(iplothelper.timestamp2str(self.end_timestamp))
            sql = "WHERE timestamp <= %i ORDER BY timestamp DESC LIMIT %i"%(self.end_timestamp,96)
        elif event.key == 'right':
            self.start_timestamp += 86400
            self.end_timestamp += 86400
            print(iplothelper.timestamp2str(self.start_timestamp))
            sql = "WHERE timestamp >= %i ORDER BY timestamp ASC LIMIT %i" %(self.start_timestamp,96)
        # 缩小

        elif event.key == 'down':
            self.start_timestamp -= 86400
            self.end_timestamp += 86400
            sql = "WHERE timestamp >= %i AND timestamp <= %i" %(self.start_timestamp,self.end_timestamp)
        else:
            return
        cObj = iplothelper.getdata_by_sql(sql)
        
        self.draw_candlestick(cObj)


    # 先不用滚轮，用左右方向键
    # def onscroll(self, event):
    #     print("%s %s" % (event.button, event.step))
    #     setp = int(event.step)
    #     print (setp)


    #     # 默认往后滚动
    #     sql = ""
    #     if setp < 0:# 往后滚动
    #         print("后 右")
    #         self.start_timestamp += 86400
    #         self.end_timestamp += 86400
    #         print(timestamp2str(self.start_timestamp))
    #         sql = "WHERE timestamp >= %i ORDER BY timestamp ASC LIMIT %i" %(self.start_timestamp,96)
    #     else:# 往前滚动
    #         print("前 左")
    #         self.end_timestamp -= 86400
    #         self.start_timestamp -= 86400
    #         print(iplot.timestamp2str(self.end_timestamp))
    #         sql = "WHERE timestamp <= %i ORDER BY timestamp DESC LIMIT %i"%(self.end_timestamp,96)

    #     cObj = iplot.getdata_by_sql(sql)
        
    #     self.draw_candlestick(cObj)

    def update(self):
        # self.im.set_data(self.X[:, :, self.ind])
        # ax.set_ylabel('slice %s' % self.ind)

        self.fig.canvas.draw()


class USALine(object):
    """美国线
    美国线又称柱线图，其构造较K线简单。
    美国线的直线部分，表示了当天行情的最高价与最低价间的波动幅度。
    左侧横线代表开盘价，右侧横线代表收盘价，绘制美国线比绘制K线简便得多。
    """

    def __init__(self, ax,fig, start_timestamp,end_timestamp):
        self.ax = ax
        self.start_timestamp =  start_timestamp
        self.end_timestamp = end_timestamp
        # 设置两边的时间戳
        
